#include <string.h>
#include <stdio.h>
#include <stdlib.h>


#include "sm2_ka.h"
#include "sm3_with_preprocess.h"

#include "sm2.h"

// 用户A私钥,用户A临时私钥,用户A临时公钥,关于用户A的杂凑值ZA,用户B公钥,用户B临时公钥,关于用户B的杂凑值ZB
//  密钥交换参数结构体
typedef struct
{
    unsigned char isSender;           // 0:响应方(用户B) 1:发起方(用户A)
    unsigned char *pbThisPriKey;      // DA,用户A私钥
    unsigned char *pbThisTempPriKey;  // RDA，用户A临时私钥
    unsigned char *pbThisTempPubKey;  // RPA，用户A临时公钥
    unsigned char *pbThisZ;           // ZA，关于用户A的杂凑值
    unsigned char *pbOtherPubKey;     // PB，用户B公钥
    unsigned char *pbOtherTempPubKey; // RPB，用户B临时公钥
    unsigned char *pbOtherZ;          // ZB，关于用户B的杂凑值
} SM2KeyExchangeParam_t;
// 密钥交换输出结构体
typedef struct
{
    unsigned char *pbSymKey;          // 输出对称秘钥
    unsigned char *pbThisVerifyData;  // 输出校验值S1
    unsigned char *pbOtherVerifyData; // 输出校验值SA
} SM2KeyExchangeOut_t;

// 使用 u8 类型的私钥获取 u8 类型的公钥
int get_pub_key_from_priv(u8 *priv_key, size_t priv_key_len, u8 **pub_key, size_t *pub_key_len)
{
    EC_KEY *key = NULL;
    const EC_GROUP *group;
    EC_POINT *pub_point = NULL;
    BIGNUM *bn_priv_key = BN_new();

    // 初始化 OpenSSL 库
    // OpenSSL_add_all_algorithms();

    // 将 u8 私钥转换为 BIGNUM
    if (!BN_bin2bn(priv_key, priv_key_len, bn_priv_key))
    {
        BN_free(bn_priv_key);
        return 0;
    }

    // 创建一个新的 EC_KEY 对象
    key = EC_KEY_new_by_curve_name(NID_sm2);
    if (!key)
    {
        BN_free(bn_priv_key);
        return 0;
    }

    // 获取椭圆曲线参数
    group = EC_KEY_get0_group(key);

    // 设置私钥
    if (!EC_KEY_set_private_key(key, bn_priv_key))
    {
        EC_KEY_free(key);
        BN_free(bn_priv_key);
        return 0;
    }

    // 计算公钥
    pub_point = EC_POINT_new(group);
    if (!pub_point || !EC_POINT_mul(group, pub_point, bn_priv_key, NULL, NULL, NULL) ||
        !EC_KEY_set_public_key(key, pub_point))
    {
        EC_KEY_free(key);
        BN_free(bn_priv_key);
        if (pub_point)
            EC_POINT_free(pub_point);
        return 0;
    }

    // 准备输出缓冲区
    *pub_key_len = i2o_ECPublicKey(key, NULL);
    if (*pub_key_len <= 0)
    {
        EC_KEY_free(key);
        BN_free(bn_priv_key);
        EC_POINT_free(pub_point);
        return 0;
    }

    *pub_key = (u8 *)malloc(*pub_key_len);
    if (!*pub_key)
    {
        EC_KEY_free(key);
        BN_free(bn_priv_key);
        EC_POINT_free(pub_point);
        return 0;
    }

    // 获取公钥
    u8 *pub_key_ptr = *pub_key;
    if (i2o_ECPublicKey(key, &pub_key_ptr) != *pub_key_len)
    {
        free(*pub_key);
        EC_KEY_free(key);
        BN_free(bn_priv_key);
        EC_POINT_free(pub_point);
        return 0;
    }

    // 清理
    EC_KEY_free(key);
    BN_free(bn_priv_key);
    EC_POINT_free(pub_point);

    // 删除所有算法
    EVP_cleanup();

    return 1;
}
void test_get_pub_key_from_priv()
{ // 示例私钥数据，这里你需要替换为你的真实私钥数据
    // 填充 priv_key_data ...
    u8 priv_key_data[32] = {0xd4, 0xde, 0x15, 0x47, 0x4d, 0xb7, 0x4d, 0x06, 0x49, 0x1c, 0x44, 0x0d, 0x30, 0x5e, 0x01, 0x24,
                            0x00, 0x99, 0x0f, 0x3e, 0x39, 0x0c, 0x7e, 0x87, 0x15, 0x3c, 0x12, 0xdb, 0x2e, 0xa6, 0x0b, 0xb3}; // SM2私钥通常为32字节

    // 用于存储公钥数据的缓冲区
    u8 *pub_key_data = NULL;
    size_t pub_key_len = 0;

    if (get_pub_key_from_priv(priv_key_data, sizeof(priv_key_data), &pub_key_data, &pub_key_len))
    {
        // 输出公钥数据
        // 注意：实际使用时应该正确处理公钥输出和内存释放
        printf("Public key: ");
        for (size_t i = 0; i < pub_key_len; ++i)
        {
            printf("%02x", pub_key_data[i]);
        }
        printf("\n");

        // 释放内存
        free(pub_key_data);
    }
    else
    {
        // 错误处理
        ERR_print_errors_fp(stderr);
    }
    return;
}
int generate_shared_key(unsigned char *shared_key, size_t shared_key_len,
                        const EC_GROUP *group, const EC_POINT *pub_key, const BIGNUM *priv_key)
{
    printf("generate_shared_key\n");
    EC_KEY *ecdh = EC_KEY_new_by_curve_name(NID_sm2);
    if (!ecdh)
        return 0; // 退出，错误处理

    // 设置 EC_KEY
    if (!EC_KEY_set_group(ecdh, group))
        goto err;

    // 设置对方的公钥
    if (!EC_KEY_set_public_key(ecdh, pub_key))
        goto err;

    // 设置自己的私钥
    if (!EC_KEY_set_private_key(ecdh, priv_key))
        goto err;

    // 计算共享密钥
    size_t key_len = ECDH_compute_key(shared_key, shared_key_len, pub_key, ecdh, NULL);
    if (key_len <= 0)
        goto err;

    // 释放资源
    EC_KEY_free(ecdh);
    return key_len;

err:
    // 错误处理
    EC_KEY_free(ecdh);
    return 0;
}

int test_generate_shared_key()
{
    printf("test_generate_shared_key\n");

    // 创建一个新的 EC_KEY 对象
    EC_KEY *key = NULL;
    key = EC_KEY_new_by_curve_name(NID_sm2);
    if (!key)
    {
        return -1;
    }
    // 获取椭圆曲线参数
    const EC_GROUP *group = EC_KEY_get0_group(key); // group跟随key被释放，不需要额外释放
    if (!group)
    {
        EC_KEY_free(key); // 获取椭圆曲线参数失败，释放 EC_KEY 对象
        return -1;
    }

    // 对方的公钥（应该是从证书或者交换过程中获得）
    u8 pub_key_data[64] = {0xc5, 0x58, 0xb4, 0x4b, 0xee, 0x53, 0x01, 0xd9, 0xf5, 0x2b, 0x44, 0xd9, 0x39, 0xbb, 0x59, 0x58, 0x4d, 0x75, 0xb9, 0x03, 0x4d, 0xd6, 0xa9, 0xfc, 0x82, 0x68, 0x72, 0x10, 0x9a, 0x65, 0x73, 0x9f, 0x32, 0x52, 0xb3, 0x5b, 0x19, 0x1d, 0x8a, 0xe0, 0x1c, 0xd1, 0x22, 0xc0, 0x25, 0x20, 0x43, 0x34, 0xc5, 0xea, 0xcf, 0x68, 0xa0, 0xcb, 0x48, 0x54, 0xc6, 0xa7, 0xd3, 0x67, 0xec, 0xad, 0x4d, 0xe7};

    BIGNUM *x = BN_bin2bn(pub_key_data, 32, NULL); // 前32字节是x坐标，后32字节是y坐标
    if (!x)
    {
        EC_KEY_free(key);
        return -1;
    }

    BIGNUM *y = BN_bin2bn(pub_key_data + 32, 32, NULL);
    if (!y)
    {
        BN_free(x);
        EC_KEY_free(key);
        return -1;
    }

    EC_POINT *pub_key = NULL;
    pub_key = EC_POINT_new(group); // 需要转换对方临时公钥到EC_POINT
    if (!pub_key)
    {
        EC_POINT_free(pub_key);
        EC_KEY_free(key);
        return -1;
    }

    // 使用 x 和 y 坐标设置 EC_POINT
    if (EC_POINT_set_affine_coordinates_GFp(group, pub_key, x, y, NULL) == 0)
    {
        printf("EC_POINT_set_affine_coordinates_GFp\n");
        EC_POINT_free(pub_key);
        EC_KEY_free(key);
        BN_free(x);
        BN_free(y);
        return -1;
    }
    // 最后，不再需要 x 和 y
    BN_free(x);
    BN_free(y);

    u8 priv_key_data[32] = {0xd4, 0xde, 0x15, 0x47, 0x4d, 0xb7, 0x4d, 0x06, 0x49, 0x1c, 0x44, 0x0d, 0x30, 0x5e, 0x01, 0x24, 0x00, 0x99, 0x0f, 0x3e, 0x39, 0x0c, 0x7e, 0x87, 0x15, 0x3c, 0x12, 0xdb, 0x2e, 0xa6, 0x0b, 0xb3}; // SM2私钥通常为32字节
    // 需要生成或者导入自己的私钥
    BIGNUM *bn_priv_key = BN_bin2bn(priv_key_data, 32, NULL); // 将 u8 私钥转换为 BIGNUM
    if (bn_priv_key == NULL)
    {
        EC_POINT_free(pub_key);
        EC_KEY_free(key);
        return 0;
    }

    // 分配共享密钥的空间
    unsigned char *shared_key = (unsigned char *)malloc(256 / 8); // 256位密钥
    if (!shared_key)
    {
        BN_free(bn_priv_key);
        EC_POINT_free(pub_key);
        EC_KEY_free(key);
        return -1;
    }
    memset(shared_key, 0, 256 / 8);

    // 生成共享密钥
    int key_len = generate_shared_key(shared_key, 256 / 8, group, pub_key, bn_priv_key);
    if (key_len <= 0)
    {
        // 错误处理
        free(shared_key);
        BN_free(bn_priv_key);
        EC_POINT_free(pub_key);
        EC_KEY_free(key);
        return -1;
    }

    // 打印共享密钥（仅用于测试）
    for (int i = 0; i < key_len; ++i)
    {
        printf("%02X", shared_key[i]);
    }
    printf("\n");

    // 清理资源
    free(shared_key);
    BN_free(bn_priv_key);
    EC_POINT_free(pub_key);
    EC_KEY_free(key);
    return 0;
}

int test_ex_sm2()
{
    unsigned char id_A[32] = {
        0x31, 0x32, 0x33, 0x34, 0x35, 0x36, 0x37, 0x38, 0x31, 0x32, 0x33, 0x34, 0x35, 0x36, 0x37, 0x38};
    unsigned char id_B[32] = {
        0x31, 0x32, 0x33, 0x34, 0x35, 0x36, 0x37, 0x38, 0x31, 0x32, 0x33, 0x34, 0x35, 0x36, 0x37, 0x38};
    unsigned char prikdata_A[32] = {
        0x81, 0xEB, 0x26, 0xE9, 0x41, 0xBB, 0x5A, 0xF1, 0x6D, 0xF1, 0x16, 0x49, 0x5F, 0x90, 0x69, 0x52,
        0x72, 0xAE, 0x2C, 0xD6, 0x3D, 0x6C, 0x4A, 0xE1, 0x67, 0x84, 0x18, 0xBE, 0x48, 0x23, 0x00, 0x29};
    unsigned char pubx_A[32] = {
        0x16, 0x0E, 0x12, 0x89, 0x7D, 0xF4, 0xED, 0xB6, 0x1D, 0xD8, 0x12, 0xFE, 0xB9, 0x67, 0x48, 0xFB,
        0xD3, 0xCC, 0xF4, 0xFF, 0xE2, 0x6A, 0xA6, 0xF6, 0xDB, 0x95, 0x40, 0xAF, 0x49, 0xC9, 0x42, 0x32};
    unsigned char puby_A[32] = {
        0x4A, 0x7D, 0xAD, 0x08, 0xBB, 0x9A, 0x45, 0x95, 0x31, 0x69, 0x4B, 0xEB, 0x20, 0xAA, 0x48, 0x9D,
        0x66, 0x49, 0x97, 0x5E, 0x1B, 0xFC, 0xF8, 0xC4, 0x74, 0x1B, 0x78, 0xB4, 0xB2, 0x23, 0x00, 0x7F};
    unsigned char prikdata_B[32] = {
        0x78, 0x51, 0x29, 0x91, 0x7D, 0x45, 0xA9, 0xEA, 0x54, 0x37, 0xA5, 0x93, 0x56, 0xB8, 0x23, 0x38,
        0xEA, 0xAD, 0xDA, 0x6C, 0xEB, 0x19, 0x90, 0x88, 0xF1, 0x4A, 0xE1, 0x0D, 0xEF, 0xA2, 0x29, 0xB5};
    unsigned char pubx_B[32] = {
        0x6A, 0xE8, 0x48, 0xC5, 0x7C, 0x53, 0xC7, 0xB1, 0xB5, 0xFA, 0x99, 0xEB, 0x22, 0x86, 0xAF, 0x07,
        0x8B, 0xA6, 0x4C, 0x64, 0x59, 0x1B, 0x8B, 0x56, 0x6F, 0x73, 0x57, 0xD5, 0x76, 0xF1, 0x6D, 0xFB};
    unsigned char puby_B[32] = {
        0xEE, 0x48, 0x9D, 0x77, 0x16, 0x21, 0xA2, 0x7B, 0x36, 0xC5, 0xC7, 0x99, 0x20, 0x62, 0xE9, 0xCD,
        0x09, 0xA9, 0x26, 0x43, 0x86, 0xF3, 0xFB, 0xEA, 0x54, 0xDF, 0xF6, 0x93, 0x05, 0x62, 0x1C, 0x4D};
    printf("=======================================\n");
    printf("密钥协商初始化：\n");
    printf("=======================================\n");

    SM2KA ka_a;
    SM2KA ka_b;
    EC_POINT *RA = NULL;
    EC_POINT *RB = NULL;
    EC_POINT *PA = NULL;
    EC_POINT *PB = NULL;

    PA = agreement_init(&ka_a, prikdata_A, pubx_A, puby_A, id_A, 16);
    dump_hex("ZA:", ka_a.Z, 16);
    PB = agreement_init(&ka_b, prikdata_B, pubx_B, puby_B, id_B, 16);
    dump_hex("ZB:", ka_b.Z, 16);
    printf("=======================================\n");
    printf("发送方A：\n");
    printf("=======================================\n");
    RA = gen_agreementdata(&ka_a);
    if (RA == NULL)
    {
        return -1;
    }
    printf("=======================================\n");
    printf("发送方B：\n");
    printf("=======================================\n");
    RB = gen_agreementdata_key(&ka_b, RA, PA, ka_a.Z);
    if (RB == NULL)
    {
        return -1;
    }
    printf("=======================================\n");
    printf("发送方A：\n");
    printf("=======================================\n");
    gen_key(&ka_a, RB, PB, ka_b.Z);

    return 0;
}
